import PIL.ImageOps
import torch
import torchvision
import time
from PIL import Image, ImageFilter
import numpy as np
import traceback
def load_data_mnist(batch_size_train, batch_size_test):
    mnist_train = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('./dataset/', train=True, download=True,
                             transform=torchvision.transforms.Compose([
                               torchvision.transforms.ToTensor(),
                               # torchvision.transforms.Normalize(
                               #   (0.1307,), (0.3081,))
                             ])),
  batch_size=batch_size_train, shuffle=True)

    mnist_test = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('./dataset/', train=False, download=True,
                             transform=torchvision.transforms.Compose([
                               torchvision.transforms.ToTensor(),
                               # torchvision.transforms.Normalize(
                               #   (0.1307,), (0.3081,))
                             ])),
  batch_size=batch_size_test, shuffle=True)
    return mnist_train, mnist_test

train_losses = []
test_acc_list = []
def train_net(net, train_iter, test_iter, optimizer, device, epochs):
    net = net.to(device)
    print("training on ", device)
    loss = torch.nn.CrossEntropyLoss()
    for epoch in range(epochs):
        train_l_sum, train_acc_sum, n, batch_count, start = 0.0, 0.0, 0, 0, time.time()
        for X, y in train_iter:
            X = X.to(device)
            y = y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            train_l_sum += l.cpu().item()
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().cpu().item()
            n += y.shape[0]
            batch_count += 1
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f, time %.1f sec'
              % (epoch + 1, train_l_sum / batch_count, train_acc_sum / n, test_acc, time.time() - start))
        train_losses.append(train_l_sum / batch_count)
        test_acc_list.append(test_acc)
        torch.save(net.state_dict(), './model.pth')
        # torch.save(optimizer.state_dict(), './optimizer.pth')
    save_list('train_losses.txt', train_losses)
    save_list('test_acc.txt', test_acc_list)

def evaluate_accuracy(data_iter, net, device=None):
    if device is None and isinstance(net, torch.nn.Module):
        # 如果没指定device就使用net的device
        device = list(net.parameters())[0].device
    acc_sum, n = 0.0, 0
    with torch.no_grad():
        for X, y in data_iter:
            if isinstance(net, torch.nn.Module):
                net.eval() # 评估模式, 这会关闭dropout
                acc_sum += (net(X.to(device)).argmax(dim=1) == y.to(device)).float().sum().cpu().item()
                net.train() # 改回训练模式
            # else:
            #     if('is_training' in net.__code__.co_varnames): # 如果有is_training这个参数
            #         # 将is_training设置成False
            #         acc_sum += (net(X, is_training=False).argmax(dim=1) == y).float().sum().item()
            #     else:
            #         acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
            n += y.shape[0]
    return acc_sum / n

def save_list(file_name, to_save):
    file = open(file_name, 'a')
    for item in to_save:
        file.write(str(item))
        file.write('\n')
    file.close()

def load_list(file_name):
    file = open(file_name, 'r')
    list = []
    while 1:
        line = file.readline().strip('\n')
        if not line:
            break
        list.append(eval(line))

    file.close()
    return list

def prepare_image(image_path):
    try:
        """
        返回处理后的图片
        image_path - 图片路径
        """
        pre = '../gui/'
        img = Image.open(pre + image_path).convert('L')

        width = float(img.size[0])
        heigth = float(img.size[1])
        print(img)
        # 创建28*28的黑色背景图片
        new_img = Image.new('L', (28, 28), (255))

        # 把图片限制在20*20像素里面（按原比例）
        if width > heigth:
            # 根据比例宽度调整高度
            nheight = int(round((20.0 / width * heigth), 0))
            # 不能让图片高为0
            if (nheight == 0):
                nheight = 1
                # 调整大小和锐化
            img = img.resize((20, nheight), Image.ANTIALIAS).filter(ImageFilter.SHARPEN)
            # 计算垂直位置
            wtop = int(round(((28 - nheight) / 2), 0))
            # 贴到背景图片上
            new_img.paste(img, (4, wtop))
        else:
            # 根据比例高度调整宽度
            nwidth = int(round((20.0 / heigth * width), 0))
            # 不能让图片宽为0
            if (nwidth == 0):
                nwidth = 1
            # 调整大小和锐化
            img = img.resize((nwidth, 20), Image.ANTIALIAS).filter(ImageFilter.SHARPEN)
            # 计算水平位置
            wleft = int(round(((28 - nwidth) / 2), 0))
            # 贴到背景图片上
            new_img.paste(img, (wleft, 4))

        # 得到像素列表
        tv = list(new_img.getdata())
        print(tv)
        # 将像素正常化为0和1。0是纯白色，1是纯黑色。
        # tva = [(255 - x) * 1.0 / 255.0 for x in tv]
        image = Image.fromarray(np.reshape(tv, (28, 28)))
        if image.mode != "RGB":
            image = image.convert('RGB')
        return image
    except Exception:
        traceback.print_exc()

def invert_image(img):
    inverted = PIL.ImageOps.invert(img)
    inverted.save('test_temp.jpg')